@contextlib.contextmanager
def execution_context(**kwargs: str) -> None:
    """Temporarily set context values for the current async task."""
    # Create a fresh copy per use
    current = {**server.execution_context_vars.get(), **kwargs}
    token = server.execution_context_vars.set(current)
    try:
        yield
    finally:
        server.execution_context_vars.reset(token)


@server.feature("initializeFunctionOptimization")
def initialize_function_optimization(params: FunctionOptimizationInitParams) -> dict[str, str]:
    with execution_context(task_id=params.task_id):
        document_uri = params.textDocument.uri
        document = server.workspace.get_text_document(document_uri)

        server.show_message_log(
            f"Initializing optimization for function: {params.functionName} in {document_uri}", "Info"
        )

        if server.optimizer is None:
            _initialize_optimizer_if_api_key_is_valid()

        server.optimizer.worktree_mode()

        original_args, _ = server.optimizer.original_args_and_test_cfg

        server.optimizer.args.function = params.functionName
        original_relative_file_path = Path(document.path).relative_to(original_args.project_root)
        server.optimizer.args.file = server.optimizer.current_worktree / original_relative_file_path
        server.optimizer.args.previous_checkpoint_functions = False

def abort_if_cancelled(cancel_event: threading.Event) -> None:
    if cancel_event.is_set():
        raise RuntimeError("cancelled")


def sync_perform_optimization(params) -> dict[str, str]:  # noqa
    server = CodeflashServerSingleton.get()
    cancel_event = server.cancel_event
    server.show_message_log(f"Starting optimization for function: {params.functionName}", "Info")

class CodeflashServerSingleton:
    _instance: CodeflashLanguageServer | None = None

    @classmethod
    def get(cls) -> CodeflashLanguageServer:
        if cls._instance is None:
            cls._instance = CodeflashLanguageServer(
                "codeflash-language-server", "v1.0", protocol_cls=CodeflashLanguageServerProtocol
            )
        return cls._instance

    def __init__(self) -> None:
        # This is a singleton class, so we don't want to initialize.
        ...